import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import torch
import os
from PIL import Image


def plot_training_curves(train_losses, val_losses, train_dice, val_dice, save_path):
    """绘制训练曲线"""
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
    
    # 绘制损失曲线
    epochs = range(1, len(train_losses) + 1)
    ax1.plot(epochs, train_losses, 'b-', label='Training Loss')
    ax1.plot(epochs, val_losses, 'r-', label='Validation Loss')
    ax1.set_title('Training and Validation Loss')
    ax1.set_xlabel('Epochs')
    ax1.set_ylabel('Loss')
    ax1.legend()
    ax1.grid(True)
    
    # 绘制Dice曲线
    ax2.plot(epochs, train_dice, 'b-', label='Training mean Dice')
    ax2.plot(epochs, val_dice, 'r-', label='Validation meanDice')
    ax2.set_title('Training and Validation mean Dice Score')
    ax2.set_xlabel('Epochs')
    ax2.set_ylabel('mean Dice Score')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()


def plot_confusion_matrix(cm, class_names, save_path, title='Confusion Matrix'):
    """绘制混淆矩阵热力图"""
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.title(title)
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()


def visualize_prediction(image, true_mask, pred_mask, save_path, class_colors_list=None, class_names=None):
    if class_colors_list is None:
        # 默认颜色映射，如果未提供则使用随机颜色
        num_classes = max(np.max(true_mask), np.max(pred_mask)) + 1 if true_mask is not None else np.max(pred_mask) + 1
        class_colors = {i: list(np.random.randint(0, 256, 3)) for i in range(num_classes)}
    else:
        class_colors = {i: color for i, color in enumerate(class_colors_list)}
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    
    # 原图
    if isinstance(image, torch.Tensor):
        if image.dim() == 3 and image.shape[0] == 3:  # CHW格式
            image_np = image.permute(1, 2, 0).cpu().numpy()
        else:
            image_np = image.cpu().numpy()
        
        # 反归一化（假设使用ImageNet标准化）
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        image_np = std * image_np + mean
        image_np = np.clip(image_np, 0, 1)
    else:
        image_np = np.array(image)
    
    axes[0].imshow(image_np)
    axes[0].set_title('Original Image')
    axes[0].axis('off')
    
    # 真实标签
    if true_mask is not None:
        if isinstance(true_mask, torch.Tensor):
            true_mask_np = true_mask.cpu().numpy()
        else:
            true_mask_np = np.array(true_mask)
        
        true_colored = create_colored_mask(true_mask_np, class_colors)
        axes[1].imshow(true_colored)
        axes[1].set_title('Ground Truth')
        axes[1].axis('off')
    else:
        axes[1].set_title('No Ground Truth')
        axes[1].axis('off')
    
    # 预测结果
    if isinstance(pred_mask, torch.Tensor):
        pred_mask_np = pred_mask.cpu().numpy()
    else:
        pred_mask_np = np.array(pred_mask)
    
    pred_colored = create_colored_mask(pred_mask_np, class_colors)
    axes[2].imshow(pred_colored)
    axes[2].set_title('Prediction')
    axes[2].axis('off')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()


def create_colored_mask(mask, class_colors):
    """将分割掩码转换为彩色图像"""
    
    h, w = mask.shape
    colored_mask = np.zeros((h, w, 3), dtype=np.uint8)
    
    for class_id, color in class_colors.items():
        colored_mask[mask == class_id] = color
    
    return colored_mask


def plot_metrics_comparison(metrics_dict, save_path, title='Metrics Comparison', class_names=None):
    """绘制不同类别的指标对比图"""
    
    metrics_names = ['dice', 'iou', 'precision', 'recall', 'f1_score']
    num_classes = len(metrics_dict['dice'])
    
    if class_names is None:
        class_names = [f'Class {j}' for j in range(num_classes)]

    fig, axes = plt.subplots(1, len(metrics_names), figsize=(20, 4))
    
    for i, metric_name in enumerate(metrics_names):
        if metric_name in metrics_dict:
            values = metrics_dict[metric_name]
            
            bars = axes[i].bar(class_names, values)
            axes[i].set_title(f'{metric_name.upper()}')
            axes[i].set_ylabel('Score')
            axes[i].set_ylim(0, 1)
            axes[i].tick_params(axis='x', rotation=45)
            
            # 在柱状图上添加数值
            for bar, value in zip(bars, values):
                axes[i].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                           f'{value:.3f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()


def save_prediction_overlay(image, pred_mask, save_path, alpha=0.5, class_colors_list=None):
    if class_colors_list is None:
        num_classes = np.max(pred_mask) + 1
        class_colors = {i: list(np.random.randint(0, 256, 3)) + [128] for i in range(num_classes)}
    else:
        class_colors = {i: color + [128] if len(color) == 3 else color for i, color in enumerate(class_colors_list)}

    # 处理图像
    if isinstance(image, torch.Tensor):
        if image.dim() == 3 and image.shape[0] == 3:
            image_np = image.permute(1, 2, 0).cpu().numpy()
        else:
            image_np = image.cpu().numpy()
        
        # 反归一化
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        image_np = std * image_np + mean
        image_np = np.clip(image_np * 255, 0, 255).astype(np.uint8)
    else:
        image_np = np.array(image)
    
    # 处理预测掩码
    if isinstance(pred_mask, torch.Tensor):
        pred_mask_np = pred_mask.cpu().numpy()
    else:
        pred_mask_np = np.array(pred_mask)
    
    # 创建叠加图像
    h, w = pred_mask_np.shape
    overlay = np.zeros((h, w, 4), dtype=np.uint8)
    
    for class_id, color in class_colors.items():
        mask_class = pred_mask_np == class_id
        overlay[mask_class] = color
    
    # 转换为PIL图像并保存
    base_image = Image.fromarray(image_np)
    overlay_image = Image.fromarray(overlay, 'RGBA')
    
    result = Image.alpha_composite(base_image.convert('RGBA'), overlay_image)
    result.save(save_path)


if __name__ == "__main__":
    # 测试绘图功能
    import tempfile
    
    # 创建临时目录
    temp_dir = tempfile.mkdtemp()
    
    # 测试训练曲线
    train_losses = [0.8, 0.6, 0.4, 0.3, 0.2]
    val_losses = [0.9, 0.7, 0.5, 0.4, 0.3]
    train_dice = [0.2, 0.4, 0.6, 0.7, 0.8]
    val_dice = [0.1, 0.3, 0.5, 0.6, 0.7]
    
    plot_training_curves(train_losses, val_losses, train_dice, val_dice, 
                        os.path.join(temp_dir, 'training_curves.png'))
    print(f"训练曲线已保存到: {os.path.join(temp_dir, 'training_curves.png')}")
    
    # 测试混淆矩阵
    cm = np.array([[50, 2, 3], [1, 45, 4], [2, 3, 48]])
    class_names = ['Background', 'Class 1', 'Class 2']
    plot_confusion_matrix(cm, class_names, os.path.join(temp_dir, 'confusion_matrix.png'))
    print(f"混淆矩阵已保存到: {os.path.join(temp_dir, 'confusion_matrix.png')}")
    
    # 测试指标对比
    metrics_dict = {
        'dice': [0.95, 0.85, 0.75],
        'iou': [0.90, 0.80, 0.70],
        'precision': [0.92, 0.88, 0.78],
        'recall': [0.94, 0.82, 0.72],
        'f1_score': [0.93, 0.85, 0.75]
    }
    plot_metrics_comparison(metrics_dict, os.path.join(temp_dir, 'metrics_comparison.png'), class_names=class_names)
    print(f"指标对比图已保存到: {os.path.join(temp_dir, 'metrics_comparison.png')}")
    
    print(f"所有测试图像保存在: {temp_dir}")


